#!/usr/bin/env python3

import threading
import time

import rclpy
from rclpy.node import Node
from rclpy.lifecycle import (LifecycleNode, LifecycleState, TransitionCallbackReturn)
from rclpy.action import ActionServer
from rclpy.action.server import ServerGoalHandle, GoalResponse, CancelResponse
from rclpy.executors import MultiThreadedExecutor
from rclpy.callback_groups import ReentrantCallbackGroup as Reen
from ros2_py_pkg.action import MoveRobot

from std_msgs.msg import Int64
from std_msgs.msg import Int64MultiArray



DISTANCE_THRESHOLD = 0.01 # meters

class MoveRobotActionServer(LifecycleNode):
    def __init__(self):
        super().__init__("move_robot_server_node")
        self.goal_lock_ = threading.Lock()
        self.goal_handle_: ServerGoalHandle = None
        self.server_activate = False
        
        self.robot_current_position_ = 50    # default position is 50 meters
        self.velocity = 1     # test velocity

        self.get_logger().info("Robot init position: " + str(self.robot_current_position_))

    def on_configure(self, state: LifecycleState):
        self.declare_parameter("robot_name", rclpy.Parameter.Type.STRING)
        self.robot_name_ = self.get_parameter("robot_name").value

        self.move_robot_action_server_ = ActionServer(
            self,
            MoveRobot,
            "move_robot_" + self.robot_name_,
            execute_callback=self.move_robot_callback,
            callback_group=Reen(),
            goal_callback=self.goal_callback,
            cancel_callback=self.cancel_callback
        )
        
        # Subscribe robot_current_position topic
        self.sub_1_ = self.create_subscription(
            Int64, 'robot_current_position', self.update_robot_position, 10)
        self.sub_2_ = self.create_subscription(
            Int64MultiArray, 'robot_current_position_velocity', 
            self.update_robot_pos_vel, 10)
        
        self.get_logger().info("Action server has been started")

        return TransitionCallbackReturn.SUCCESS
    
    def on_activate(self, state: LifecycleState):
        self.get_logger().info("Activate node")
        self.server_activate = True
        return super().on_activate(state)
    
    def on_deactivate(self, state: LifecycleState):
        self.get_logger().info("Deactivate node")
        self.server_activate = False
        with self.goal_lock_:
            if self.goal_handle_ is not None and self.goal_handle_.is_active:
                self.goal_handle_.abort()
        return super().on_deactivate(state)
    
    def on_cleanup(self, state: LifecycleState):
        self.get_logger().info("Clean up server")
        self.undeclare_parameter("robot_name")
        self.robot_name_ = ""
        self.move_robot_action_server_.destroy()
        return TransitionCallbackReturn.SUCCESS
    
    def on_shutdown(self, state: LifecycleState):
        self.get_logger().info("Shut down server")
        self.move_robot_action_server_.destroy()
        return TransitionCallbackReturn.SUCCESS
    
    def on_error(self, state: LifecycleState):
        self.move_robot_action_server_.destroy()
        return TransitionCallbackReturn.FAILURE

    def move_robot_callback(self, goal_handle: ServerGoalHandle):
        with self.goal_lock_:
            self.goal_handle_ = goal_handle

        self.get_logger().info("Execute Goal")
        start_time = self.get_clock().now().to_msg().sec
        goal_position = goal_handle.request.position
        velocity = goal_handle.request.velocity

        result = MoveRobot.Result()
        feedback_msg = MoveRobot.Feedback()

        self.get_logger().info("Goal Position: " + str(goal_position))
        while rclpy.ok():
            if not goal_handle.is_active:
                result.position = self.robot_current_position_
                result.message = "Preempted by another goal"
                return result
            
            if goal_handle.is_cancel_requested:
                result.position = self.robot_current_position_
                if goal_position == self.robot_current_position_:
                    result.message = "Success after cancel request"
                    goal_handle.succeed()
                else:
                    result.message = "Canceled"
                    goal_handle.canceled()
                return result

            distance_to_goal = goal_position - self.robot_current_position_

            if abs(distance_to_goal) <= DISTANCE_THRESHOLD:
                result.position = distance_to_goal
                result.message = "Success"
                goal_handle.succeed()
                return result
            elif distance_to_goal > DISTANCE_THRESHOLD:
                if distance_to_goal >= velocity:
                    self.robot_current_position_ += velocity
                else:
                    self.robot_current_position_ += distance_to_goal
            else:
                if abs(distance_to_goal) >= velocity:
                    self.robot_current_position_ -= velocity
                else:
                    self.robot_current_position_ -= abs(distance_to_goal)

            self.get_logger().info("Robot current position: " + str(self.robot_current_position_))
            feedback_msg.current_position = self.robot_current_position_
            goal_handle.publish_feedback(feedback_msg)
            # rclpy.spin_once(self, timeout_sec=1)
            time.sleep(1.0)

    def goal_callback(self, goal_request: MoveRobot.Goal):
        self.get_logger().info("Received a new goal")
        
        if not self.server_activate:
            self.get_logger().warn("Server not acetivated yet")
            return GoalResponse.REJECT
        
        if goal_request.position not in range(0, 100) or goal_request.velocity <= 0:
            self.get_logger().warn("Invalid position/velocity, reject goal")
            return GoalResponse.REJECT
        
        # New goal is valid, abort previous goal and accept new goal
        if self.goal_handle_ is not None and self.goal_handle_.is_active:
            self.get_logger().info("New goal is valid, abort previous goal and accept new goal")
            self.goal_handle_.abort()
        
        self.get_logger().info("Accept goal")
        return GoalResponse.ACCEPT

    def cancel_callback(self, goal_handle: ServerGoalHandle):
        self.get_logger().info("Received a cancel request")

        return CancelResponse.ACCEPT

    def update_robot_position(self, postition: Int64):
        self.robot_current_position_ = postition.data
    
    def update_robot_pos_vel(self, goal: Int64MultiArray):
        self.robot_current_position_ = goal.data[0]
        self.velocity = goal.data[1]


def main(argc=None):
    rclpy.init()
    move_robot_server_node = MoveRobotActionServer()

    try:
        while rclpy.ok:
            rclpy.spin(move_robot_server_node, executor=MultiThreadedExecutor())
    except KeyboardInterrupt:
        print("Terminating Node...")
        move_robot_server_node.move_robot_action_server_.destroy()
        move_robot_server_node.destroy_node()

if __name__ == "__main__":
    main()